import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

# 加载 digits 数据集
digits = load_digits()
X, y = digits.data, digits.target

# 划分数据集为训练集和测试集
# 测试集占比30%，训练集占70%，随机种子设为42（确保结果可复现）
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.5,  # 测试集比例（也可用train_size指定训练集比例）
    random_state=42,  # 随机种子（任意整数，固定后划分结果不变）
    shuffle=True  # 默认True，随机打乱数据；False则按原顺序取后30%为测试集
)

# 查看划分结果
print(f"训练集样本数：{X_train.shape[0]}")
print(f"测试集样本数：{X_test.shape[0]}")

# 训练 KNN 分类器
knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_train, y_train)
# 选择几个测试样本
sample_indices = [0, 1, 2, 3, 4]
sample_data = X_test[sample_indices]

# 使用训练好的模型进行预测
predictions = knn.predict(sample_data)

# 绘制样本图像
plt.figure(figsize=(10, 4))
for i, index in enumerate(sample_indices):
    plt.subplot(1, len(sample_indices), i + 1)
    plt.imshow(X_test[index].reshape(8, 8), cmap='gray')
    plt.title(f'Pred: {predictions[i]}')
    plt.axis('off')

plt.show()
